from PyQt5.QtWidgets import *
from PyQt5.QtCore import Qt, QBasicTimer
from PyQt5 import QtGui
from PyQt5 import QtWidgets
from PyQt5.QtCore import QObject,pyqtSignal
from UI import Ui_MainWindow
from codes.test_ClassSR import ClassSR
from codes.test_srrescycgan import SRResCycGan
from threading import Thread
from datetime import datetime

import sys
import os
from os import path as osp


def get_timestamp():
    return (datetime.now().strftime('%y%m%d-%H:%M:%S'))


class SignalStore(QObject):
    progress_update = pyqtSignal(int)
    # 还可以定义其他作用的信号
    pic_output = pyqtSignal()
    ssim_output = pyqtSignal(str)
    psnr_output = pyqtSignal(str)
    select_ban_sig = pyqtSignal()
    select_allow_sig = pyqtSignal()



class MainForm(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super(MainForm, self).__init__()
        self.setupUi(self)

        root_path = osp.abspath(osp.join(__file__, osp.pardir))
        self.tmp_result_path = osp.join(root_path, 'temp', 'result')

        so.progress_update.connect(self.setProgress)
        so.pic_output.connect(self.finish)
        so.psnr_output.connect(self.set_psnr)
        so.ssim_output.connect(self.set_ssim)
        so.select_ban_sig.connect(self.select_ban)
        so.select_allow_sig.connect(self.select_allow)


        self.inputPath = ''       #输入图片路径
        self.outputdir = ''     #输出图片路径

        self.pic = QtGui.QPixmap()

        self.output_mode = 0  #   0：输出   1：评测
        self.model_type = 1   #   1：模型1（ClassSR+FSRCNN)    2：模型2(ClassSR+RCAN)     3：模型3(ResCycGan)

        self.button_group_1 = QButtonGroup(self)
        self.button_group_2 = QButtonGroup(self)
        self.button_group_1.addButton(self.mode_1, 11)
        self.button_group_1.addButton(self.mode_2, 12)
        self.button_group_2.addButton(self.model_1, 21)
        self.button_group_2.addButton(self.model_2, 22)
        self.button_group_2.addButton(self.model_3, 23)

        self.mode_1.setChecked(True)
        self.model_1.setChecked(True)

        self.progressBar.setRange(0,10)

        self.button_group_1.buttonClicked.connect(self.change_mode)
        self.button_group_2.buttonClicked.connect(self.change_model)
        self.add_btn.clicked.connect(self.open_file)
        self.output_path_btn.clicked.connect(self.save_file)
        self.start_btn.clicked.connect(self.start)
        self.open_dir_btn.clicked.connect(self.open_dir)

        # 是否正在进行处理
        self.ongoing = False

    def select_ban(self):
        self.model_1.setEnabled(False)
        self.model_2.setEnabled(False)
        self.model_3.setEnabled(False)
        self.mode_1.setEnabled(False)
        self.mode_2.setEnabled(False)

    def select_allow(self):
        self.model_1.setEnabled(True)
        self.model_2.setEnabled(True)
        self.model_3.setEnabled(True)
        self.mode_1.setEnabled(True)
        self.mode_2.setEnabled(True)

    def open_dir(self):
        if self.outputdir:
            os.startfile(self.outputdir)
        else:
            QMessageBox.warning(self, '警告','未选择输出路径')



    def setProgress(self,value):
        self.progressBar.setValue(value)

    # 选择输入图片
    def open_file(self):
        if self.ongoing:
            QMessageBox.critical (self, '错误','任务正在进行中')
        else:
            self.inputPath = ''
            self.input_pic.setText("请选择输入图片...")
            self.output_pic.setText("输出图片")
            self.input_path_txt.setText("")
            filename = QtWidgets.QFileDialog.getOpenFileName(self, "选取图片", os.getcwd(),
                                                                      "*.jpg;*.png;;All Files(*)")
            self.inputPath = filename[0]

            if self.inputPath:
                self.pic = QtGui.QPixmap(self.inputPath)
                height = self.pic.height()
                width = self.pic.width()
                if width > self.input_pic.width():
                    width = self.input_pic.width()
                if height > self.input_pic.height():
                    height = self.input_pic.height()
                self.pic = self.pic.scaled(width, height, Qt.KeepAspectRatio,
                                                    Qt.SmoothTransformation)
                self.input_path_txt.setText(os.path.basename(self.inputPath))
                self.input_pic.setPixmap(self.pic)
                so.psnr_output.emit("")
                so.ssim_output.emit("")
            else:
                print('open err')

    # 选择保存路径
    def save_file(self):
        if self.ongoing:
            QMessageBox.critical (self, '错误','任务正在进行中')
        else:
            self.outputdir = ''
            self.output_pic.setText("输出图片")
            self.output_path_txt.setText("请选择输出路径...")

            filename = QtWidgets.QFileDialog.getExistingDirectory(self, "选择文件夹", "./")
            if filename != '':
                self.outputdir = filename
                self.output_path_txt.setText(self.outputdir)


    # 开始运行
    def start(self):
        if self.ongoing:
            QMessageBox.critical (self, '错误','任务正在进行中')
        else:
            if not self.inputPath:
                print('input err')
                QMessageBox.warning(self, '警告','请选择输入图片')
            elif not self.outputdir and not self.output_mode:
                print('output err')
                QMessageBox.warning(self,'警告', '请选择输出路径')

            else:
                # self.onstart()
                print('start')
                '''
                在这调用接口，传递图片路径:self.inputPath, model_mode
                '''
                self.output_pic.setText("输出图片")
                so.psnr_output.emit("")
                so.ssim_output.emit("")
                if self.model_type == 1 or self.model_type == 2:
                    is_fast = (self.model_type == 1)
                    worker = Thread(target=ClassSR, args=(so, self.inputPath, self.outputdir, self.output_mode, is_fast))
                else:
                    worker = Thread(target=SRResCycGan,args=(so, self.inputPath, self.outputdir, self.output_mode))
                so.progress_update.emit(1)
                self.ongoing = True
                print(get_timestamp())
                worker.start()
                so.select_ban_sig.emit()



    def finish(self):
        print(get_timestamp())
        self.ongoing = False
        so.select_allow_sig.emit()
        if not self.output_mode:
            outputPath = os.path.join(self.outputdir,os.path.basename(self.inputPath))
        else:
            outputPath = os.path.join(self.tmp_result_path,os.path.basename(self.inputPath))
        print(outputPath)
        output = QtGui.QPixmap(outputPath)
        height = output.height()
        width = output.width()
        if width > self.output_pic.width():
            width = self.output_pic.width()
        if height > self.output_pic.height():
            height = self.output_pic.height()
        output = output.scaled(width, height, Qt.KeepAspectRatio,
                               Qt.SmoothTransformation)
        self.output_pic.setPixmap(output)


    def set_psnr(self,str):
        self.psnr.setText(str)

    def set_ssim(self, str):
        self.ssim.setText(str)

    def change_mode(self):
        so.ssim_output.emit("")
        so.psnr_output.emit("")
        if self.mode_1.isChecked():
            self.output_mode = 0
            self.output_path_btn.setEnabled(True)
            self.open_dir_btn.setEnabled(True)
            self.verticalFrame_2.setStyleSheet("* { background-color:gray }")
        elif self.mode_2.isChecked():
            self.output_mode = 1
            self.output_path_btn.setEnabled(False)
            self.open_dir_btn.setEnabled(False)
            self.verticalFrame_2.setStyleSheet("* { background-color:white }")
        print(self.output_mode)

    # 切换模型
    def change_model(self):
        so.ssim_output.emit("")
        so.psnr_output.emit("")
        if self.model_1.isChecked():
            self.model_type = 1
            # self.mode_2.setEnabled(True)
        elif self.model_2.isChecked():
            self.model_type = 2
            # self.mode_2.setEnabled(True)
        elif self.model_3.isChecked():
            self.model_type = 3
            # self.mode_2.setEnabled(False)

        print(self.model_type)



if __name__ == "__main__":
    so = SignalStore()
    app = QtWidgets.QApplication(sys.argv)
    app.setWindowIcon(QtGui.QIcon("logo.png"))
    win = MainForm()
    win.show()
    sys.exit(app.exec_())
